import pandas as pd
from sqlalchemy import text
from tqdm import tqdm
from conn import clip_annotated_record, mysql_engine, clip_col
from loguru import logger
from pydantic import BaseModel, Field
from itertools import batched

# 对应的标注类型

OCC = (21,)
LANE = (2,)
GOP = (23,)
D_PVB = (15, 16)  # 行车
P_PVB = (17,)  # 泊车


# 供应商映射
SUPPLIER_DICT = {
    0: "智目",
    1: "恺望",
    2: "百度",
    3: "ZDriveAI",
    4: "文远知行",
    5: "腾讯",
    6: "柏川",
    7: "曼孚",
    8: "百度4D",
    9: "整数",
    11: "澳鹏",
    10: "数据堂",
    12: "标贝",
    13: "ZdriveLabel",
    14: "曼孚-2",
}

# 同源查询的总体思路
# 1. 从数据库中查询到每个标注类型的 clip 列表(如何持久话存储，用什么存储更好)
# 2. 根据需求做并集或差集处理


class Clip(BaseModel):
    clip_id: str = Field(..., title="Clip ID")
    clip_name: str | None = Field(None, title="Clip Name")
    annotate_type: int = Field(..., title="标注类型")
    annotate_name: str = Field(..., title="标注类型名称")
    annotate_id: str = Field(..., title="标注 ID")
    supplier: str = Field(..., title="供应商")
    state: int | None = Field(None, title="状态")
    collect_at: str | None = Field(None, title="采集时间")


def convert2df(dashboard_list: list[Clip]):
    # field_titles = {name: info.title for name, info in Clip.model_fields.items()}
    dashboard_json_list = [dashboard.model_dump() for dashboard in dashboard_list]
    df = pd.DataFrame(dashboard_json_list)
    # df.rename(columns=field_titles, inplace=True)
    return df


def df2excel(df: pd.DataFrame, excel_path: str):
    with pd.ExcelWriter(excel_path) as writer:
        df.to_excel(writer, index=False)


def curate_query(route_tag: str):
    def _label_clips_query(anno_types: tuple[int], anno_name: str):
        clips = []
        with mysql_engine.connect() as conn:
            sql = text(
                """
                select id, supplier, annotate_id, annotate_type
                from send_label_task_tab 
                where annotate_type in :anno_type_ids
                """
            ).bindparams(anno_type_ids=anno_types)

            all_task = conn.execute(sql).fetchall()
            for (
                task_id,
                supplier,
                annotate_id,
                annotate_type,
            ) in tqdm(all_task):
                sql = text(
                    """
                    select clip_id
                    from send_label_record_tab 
                    where send_label_task_id = :task_id and data_type = 1
                    group by clip_id
                    """
                ).bindparams(task_id=task_id)
                cur_task_clips = conn.execute(sql).fetchall()
                for (clip_id,) in cur_task_clips:
                    clips.append(
                        Clip(
                            clip_id=clip_id,
                            annotate_type=annotate_type,
                            annotate_name=anno_name,
                            annotate_id=annotate_id,
                            supplier=SUPPLIER_DICT[int(supplier)],
                        )
                    )
        clip_ids = [clip.clip_id for clip in clips]
        clip_id_map = {clip.clip_id: clip for clip in clips}

        valid_tags_clips = []
        # NOTE clip 数量可能很多，只能分批加载数据
        for batch_clip_ids in tqdm(
            batched(clip_ids, 1000), total=len(clip_ids) // 1000
        ):
            batch_clip_infos = clip_col.find(
                {"_id": {"$in": batch_clip_ids}},
                {"name": 1, "tags": 1, "collect_at": 1},
            )
            for clip_info in batch_clip_infos:
                clip_id = clip_info["_id"]
                clip = clip_id_map.get(clip_id)
                clip.clip_name = clip_info["name"]
                clip.collect_at = clip_info["collect_at"]

                if (
                    route_tag in clip_info["tags"].get("route", [])
                    and clip_info["collect_at"] >= "2024-11-01"
                ):
                    valid_tags_clips.append(clip)

            batch_clip_infos = clip_annotated_record.find(
                {"clip_id": {"$in": batch_clip_ids}},
                {"name": 1, "state": 1, "clip_id": 1},
            )
            for clip_info in batch_clip_infos:
                clip_id = clip_info["clip_id"]
                clip = clip_id_map.get(clip_id)
                clip.state = clip_info["state"]

        return valid_tags_clips

    lane_clips = _label_clips_query(OCC, "OCC")
    lane_clip_id_set = set([clip.clip_id for clip in lane_clips])
    lane_clip_map = {clip.clip_id: clip for clip in lane_clips}
    logger.info(f"lane_clip_ids: {len(lane_clips)}")

    gop_clips = _label_clips_query(GOP, "GOP")
    gop_clip_id_set = set([clip.clip_id for clip in gop_clips])
    gop_clip_map = {clip.clip_id: clip for clip in gop_clips}
    logger.info(f"gop_clip_ids: {len(gop_clips)}")

    pvb_clips = _label_clips_query(P_PVB, "PVB")
    pvb_clip_id_set = set([clip.clip_id for clip in pvb_clips])
    pvb_clip_map = {clip.clip_id: clip for clip in pvb_clips}
    logger.info(f"pvb_clip_ids: {len(pvb_clips)}")

    clip_id_set_map = {
        "OCC": lane_clip_id_set,
        "GOP": gop_clip_id_set,
        "PVB": pvb_clip_id_set,
    }

    COMBINATIONS = {
        "OCC&GOP-PVB": (
            clip_id_set_map["OCC"] & clip_id_set_map["GOP"] - clip_id_set_map["PVB"]
        ),
        "OCC&PVB-GOP": (
            clip_id_set_map["OCC"] & clip_id_set_map["PVB"] - clip_id_set_map["GOP"]
        ),
        "GOP&PVB-OCC": (
            clip_id_set_map["GOP"] & clip_id_set_map["PVB"] - clip_id_set_map["OCC"]
        ),
        "OCC-GOP-PVB": (
            clip_id_set_map["OCC"] - clip_id_set_map["GOP"] - clip_id_set_map["PVB"]
        ),
        "GOP-PVB-OCC": (
            clip_id_set_map["GOP"] - clip_id_set_map["PVB"] - clip_id_set_map["OCC"]
        ),
        "PVB-OCC-GOP": (
            clip_id_set_map["PVB"] - clip_id_set_map["OCC"] - clip_id_set_map["GOP"]
        ),
    }

    dfs = {}

    for combination, clip_id_set in COMBINATIONS.items():
        logger.info(f"{combination}: {len(clip_id_set)}")
        clips = []
        for clip_id in clip_id_set:
            if clip_id in lane_clip_map:
                clips.append(lane_clip_map[clip_id])
            if clip_id in gop_clip_map:
                clips.append(gop_clip_map[clip_id])
            if clip_id in pvb_clip_map:
                clips.append(pvb_clip_map[clip_id])
        clips_df = convert2df(clips)
        dfs[combination] = clips_df

    # 将所有 DataFrame 写入 Excel 的不同 sheet
    with pd.ExcelWriter(f"{route_tag}.xlsx") as writer:  # output.xlsx 是输出文件名
        for sheet_name, df in dfs.items():
            df.to_excel(writer, sheet_name=sheet_name, index=False)


def label_query():
    def _label_clips_query(anno_type_query):
        return clip_annotated_record.find(
            {"annotate_type": anno_type_query, "state": 2},
            {
                "clip_id": 1,
                "name": 1,
                "supplier": 1,
                "annotate_id": 1,
                "annotate_type": 1,
            },
        ).to_list()


def main():
    # curate_query("highway")
    # curate_query("city_road")
    curate_query("indoor_parking")
    curate_query("open_parking")


if __name__ == "__main__":
    main()
